multi-head masked self-attention

多头遮蔽自注意力是Transformer 模型中的一个重要组成部分，它在自然语言处理（NLP）任务中发挥着重要作用。

什么是自注意力？

自注意力，也称为自我注意力或内部注意力，是一种特殊的注意力机制，它允许模型在处理输入序列的每个元素时，关注输入序列的其他部分。在处理特定词语时，自注意力层可以查看输入句子中的其他词语，这有助于模型更好地理解和表示当前词语。

什么是多头注意力？

多头注意力是自注意力的一个扩展，它通过多个自注意力层并行处理输入，每个自注意力层都有自己的权重矩阵。这种方法可以扩展模型关注不同位置的能力，并且可以让注意力层在不同的“表示子空间”中进行操作。

什么是遮蔽自注意力？

在多头遮蔽自注意力中，每个头都有自己的查询（Q）、键（K）和值（V）矩阵，这些矩阵是通过学习得到的。对于输入序列中的每个词，我们都会生成一个查询向量、一个键向量和一个值向量。然后，我们计算查询向量与所有键向量的点积，得到一个分数，这个分数表示当前词与其他词的关联程度。接着，我们通过Softmax函数将这些分数转化为概率，然后用这些概率加权值向量，得到最后的输出。

在多头注意力中，我们并行进行这个过程，每个头得到一个输出。然后，我们将所有头的输出拼接起来，通过一个线性变换得到最终的输出。

在遮蔽自注意力中，我们在计算分数时，会将未来位置的分数设为负无穷，这样在经过softmax函数后，未来位置的概率就会接近于零，从而实现了遮蔽的效果。

实现

class CausalSelfAttention(nn.Module):
    """
    一个带有投影的普通的多头掩码自注意力层。
    可以在这里使用torch.nn.MultiheadAttention，但我在这里包含了一个显式的实现，以表明这里没有什么可怕的东西。
    """

    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        # 为所有头的键、查询、值投影，但在批处理中
        # self.c_attn变量是一个线性层，它接收嵌入维度（config.n_embd）并输出3倍的嵌入维度。
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
        # 输出投影
        # self.c_proj变量是另一个线性层，它接收嵌入维度并输出相同的嵌入维度。
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
        # 正则化
        # self.attn_dropout和self.resid_dropout变量是具有在配置中指定的丢弃概率的丢弃层。
        self.attn_dropout = nn.Dropout(config.attn_pdrop)
        self.resid_dropout = nn.Dropout(config.resid_pdrop)
        # 因果掩码，以确保注意力仅应用于输入序列的左侧
        self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
                                     .view(1, 1, config.block_size, config.block_size))
        # self.n_head和self.n_embd变量分别是头的数量和嵌入维度。
        self.n_head = config.n_head
        self.n_embd = config.n_embd

        # 其它
        # self.bias变量是一个缓冲区，包含一个因果掩码，以确保注意力只应用于输入序列的左侧。

    def forward(self, x):
        B, T, C = x.size() # 批处理大小，序列长度，嵌入维度（n_embd）

        # 为批处理中的所有头计算查询、键、值，并将头向前移动以成为批处理维度
        # 输入序列 x 经过一个线性层 self.c_attn，得到三个输出 q、k 和 v。
        q, k ,v  = self.c_attn(x).split(self.n_embd, dim=2)
        # 然后，这三个输出被分别拆分成 n_head 个头，每个头的维度为 C // n_head
        # 接着，k、q 和 v 分别被重塑为 (B, nh, T, hs) 的形状，并通过 transpose 函数将头维度移动到批处理维度之前。
        # 最后，这三个张量被返回，以便进行后续的注意力计算 
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

        # 因果自注意力；自注意力：(B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        # @ 是矩阵乘法运算符。
        # q 和 k.transpose(-2, -1) 是两个矩阵
        # 它们被乘以一个标量 (1.0 / math.sqrt(k.size(-1)))。这个标量是为了缩放矩阵乘积，以便更好地处理梯度
        # k.transpose(-2, -1) 是将张量 k 的倒数第二维和倒数第一维进行转置。
        # 在这里，k 是一个形状为 (B, nh, T, hs) 的张量，其中 B 是批处理大小，nh 是头的数量，T 是序列长度，hs 是嵌入维度除以头的数量。
        # 因此，k.transpose(-2, -1) 将 k 的最后两个维度进行转置，即将头维度和序列长度维度交换
        # 首先，通过矩阵乘法计算出注意力矩阵att，其中q、k、v分别代表查询、键、值的矩阵，(B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)。  
        # 然后，通过除以k.size(-1)的平方根来缩放注意力矩阵。
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        # 接着，将注意力矩阵中bias为0的位置用负无穷替换，以便在softmax操作中被忽略。
        att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
        # 然后，对注意力矩阵进行softmax操作，得到归一化的注意力矩阵。
        att = F.softmax(att, dim=-1)
        # 接下来，对注意力矩阵进行dropout操作。
        att = self.attn_dropout(att)
        # 最后，将注意力矩阵与值矩阵v相乘，得到输出y，(B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)。
        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        # 最后，将y进行重新组装，得到(B, T, C)的输出 
        y = y.transpose(1, 2).contiguous().view(B, T, C) # 重新组装所有头的输出并排

        # 输出投影
        # y 是输入张量，self.c_proj 是一个线性层，将 y 投影到一个新的维度上。
        # self.resid_dropout 是一个 dropout 层，用于防止过拟合。
        y = self.resid_dropout(self.c_proj(y))
        # 最后，函数返回投影后的张量 y 
        return y

本文作者：Maeiee

本文链接：multi-head masked self-attention

版权声明：如无特别声明，本文即为原创文章，版权归 Maeiee 所有，未经允许不得转载！

喜欢我文章的朋友请随缘打赏，鼓励我创作更多更好的作品！